# -*- coding: utf-8 -*-


import numpy as np
import copy


def softmax(x):
    probs = np.exp(x - np.max(x))
    probs /= np.sum(probs)
    return probs


class TreeNode(object):
    """A node in the MCTS tree.

    Each node keeps track of its own value Q, prior probability P, and
    its visit-count-adjusted prior score u.
    """

    def __init__(self, parent, _P): #prior_p_policy, prior_p_value):
        self._parent = parent
        self._children = {}  # a map from action to TreeNode
        self._n_visits = 0
        self._Q = 0
        self._u = 0
        
        # self._P_policy = prior_p_policy
        # self._P_value = prior_p_value
        self._P = _P
        
        self._reward = 0

    def expand(self, action_priors, _is_self_play = 1):
        """Expand tree by creating new children.
        action_priors: a list of tuples of actions and their prior probability
            according to the policy function.
        """
        ## addd noise here
        action_priors_copy = copy.deepcopy(action_priors)
        action_length = 32 * len(list(action_priors_copy)[:])
        dirichlet_noise = 0.25 * np.random.dirichlet(0.08*np.ones(action_length))
        tmp_position = 0 
        for action, action_prob, action_value_prob_all in action_priors:
            for action_value, action_value_prob in zip(list(range(32)), action_value_prob_all):
                if (action * 32 + action_value) not in self._children:
                    if (_is_self_play == 1):
                        predict_prior = 0.75 * action_prob * action_value_prob
                        self._children[action * 32 + action_value] = TreeNode(self, predict_prior + dirichlet_noise[tmp_position])
                        tmp_position += 1
                    else:
                        predict_prior = action_prob * action_value_prob
                        self._children[action * 32 + action_value] = TreeNode(self, predict_prior)
                        
                    
        # for action, action_prob, action_value_prob_all in action_priors:
        #     for action_value, action_value_prob in zip(list(range(32)), action_value_prob_all):
        #         if (action * 32 + action_value) not in self._children:
        #             self._children[action * 32 + action_value] = TreeNode(self, action_prob * action_value_prob)

    def output_value(self, c_1, c_2):
        
        _u = ( ( c_1 + np.log( (self._parent._n_visits + c_2 + 1) / c_2 ) )  * self._P * np.sqrt(self._parent._n_visits) / (1 + self._n_visits)) 
        return self._reward, self._Q, _u
        
    
    def select(self, c_1, c_2):
        """Select action among children that gives maximum action value Q
        plus bonus u(P).
        Return: A tuple of (action, next_node)
        """

        Q_estimate = self.get_Q_estimate()
        
        return max(self._children.items(),
                   key=lambda act_node: act_node[1].get_value(c_1, c_2, Q_estimate))

    def update(self, leaf_value):
        """Update node values from leaf evaluation.
        leaf_value: the value of subtree evaluation from the current player's
            perspective.
        """
        # Count visit.
        self._n_visits += 1
        # Update Q, a running average of values for all visits.
        self._Q += 1.0 * (leaf_value - self._Q) / self._n_visits

    def update_recursive(self, leaf_value):
        """Like a call to update(), but applied recursively for all ancestors.
        """
        # If it is not root, this node's parent should be updated first.
        if self._parent:
            self._parent.update_recursive(leaf_value + self._reward)
        self.update(leaf_value)

    def get_value(self, c_1, c_2, Q_estimate):
        """Calculate and return the value for this node.
        It is a combination of leaf evaluations Q, and this node's prior
        adjusted for its visit count, u.
        c_puct: a number in (0, inf) controlling the relative impact of
            value Q, and prior probability P, on this node's score.
        c_1:
        c_2:
        """
        # self._u = ( c_puct * self._P * np.sqrt(self._parent._n_visits) / (1 + self._n_visits))
        self._u = ( ( c_1 + np.log( (self._parent._n_visits + c_2 + 1) / c_2 ) ) * self._P * np.sqrt(self._parent._n_visits) / (1 + self._n_visits)) 
        
        if (self._n_visits == 0):
            Q = Q_estimate
        else:
            Q = self._Q
        
        return self._reward + Q + self._u
    
    # def get_Q_estimate(self,):
        
    #     if (self._parent is None):
    #         Q_estimate = 0
    #     else:
    #         Q_estimate = self.get_Q_estimate(self._parent)
            
    #     Q_normalize = 1
        
    #     for _key, _node in self._children.items():
    #         if (_node._n_visits != 0):
    #             Q_estimate += (_node._reward + _node._Q)
    #             Q_normalize += 1
        
    #     return Q_estimate / Q_normalize
    
    def get_Q_estimate(self,):
        
        if (self._parent is None):
            Q_estimate = 0
        else:
            Q_estimate = self._parent.get_Q_estimate()
            
        Q_normalize = 1
        
        for _key, _node in self._children.items():
            if (_node._n_visits != 0):
                Q_estimate += (_node._reward + _node._Q)
                Q_normalize += 1
        
        return Q_estimate / Q_normalize

    def is_leaf(self):
        """Check if leaf node (i.e. no nodes below this have been expanded)."""
        return self._children == {}

    def is_root(self):
        return self._parent is None


class MCTS_candq(object):
    """An implementation of Monte Carlo Tree Search."""

    def __init__(self, policy_value_fn, c_1, c_2, n_playout=10000, _is_self_play = 1):
        """
        policy_value_fn: a function that takes in a board state and outputs
            a list of (action, probability) tuples and also a score in [-1, 1]
            (i.e. the expected value of the end game score from the current
            player's perspective) for the current player.
        c_puct: a number in (0, inf) that controls how quickly exploration
            converges to the maximum-value policy. A higher value means
            relying on the prior more.
        """
        self._root = TreeNode(None, 1.0)
        self._policy = policy_value_fn
        self._c_1 = c_1
        self._c_2 = c_2
        self._n_playout = n_playout
        self._is_self_play = _is_self_play

    def _playout(self, state):
        """Run a single playout from the root to the leaf, getting a value at
        the leaf and propagating it back through its parents.
        State is modified in-place, so a copy must be provided.
        """
        node = self._root
        while(1):
            if node.is_leaf():
                break
            # Greedily select next move.
            action, node = node.select(self._c_1, self._c_2)
            #here
            # print(action, node.output_value(self._c_1, self._c_2))
            # input()
            action_policy, action_policy_value = action // 32, action % 32
            state.do_move(action_policy, action_policy_value)
            node._reward = state.gflops_list[-1]

        # Evaluate the leaf using a network which outputs a list of
        # (action, probability) tuples p and also a score v in [-1, 1]
        # for the current player.

        action_probs, leaf_value = self._policy(state)
        
        #here for print prob
        # if (node == self._root):
        #     action_probs_visual = copy.deepcopy(action_probs)
        #     print(list(action_probs_visual))
        #     input()
        
        # policy_prob, policy_value, leaf_value = self._policy(state)
        # policy_prob, policy_value, value = ddp_model(feature_buffer, coordinate_buffer, density_buffer)
        
        leaf_value = leaf_value.cpu().numpy() # new change for this 
        # Check for end of game.
        end = state.game_end()
        if not end:
            node.expand(action_probs, _is_self_play = self._is_self_play)
        else:
            # for end state  return the "true" leaf_value
            # if winner == -1:  # tie
            #     leaf_value = 0.0
            # else:
            #     leaf_value = (
            #         1.0 if winner == state.get_current_player() else -1.0
            #     )
            # leaf_value = state.gflops_list[-1]
            leaf_value = 0
        
        # Update value and visit count of nodes in this traversal.
        node.update_recursive(leaf_value)

    def get_move_probs(self, state, temp=1e-3):
        """Run all playouts sequentially and return the available actions and
        their corresponding probabilities.
        state: the current game state
        temp: temperature parameter in (0, 1] controls the level of exploration
        """
        # print(self._n_playout)
        for n in range(self._n_playout):
            state_copy = copy.deepcopy(state)
            self._playout(state_copy)
            
        
        # calc the move probabilities based on visit counts at the root node
        act_visits = [(act, node._n_visits)
                      for act, node in self._root._children.items()]
        
        acts, visits = zip(*act_visits)
        
        # print(acts)
        # print(visits)
        
        # acts_move = acts // 32
        # acts_move_value = acts % 32
        
        act_probs = softmax( 1.0 / temp * np.log(np.array(visits) + 1e-10) )

        return acts, act_probs

    def update_with_move(self, last_move):
        """Step forward in the tree, keeping everything we already know
        about the subtree.
        """
        if last_move in self._root._children:
            self._root = self._root._children[last_move]
            self._root._parent = None
        else:
            self._root = TreeNode(None, 1.0)

    def __str__(self):
        return "MCTS"


class MCTSPlayer_candq(object):
    """AI player based on MCTS"""

    def __init__(self, policy_value_function,
                 c_1, c_2, n_playout=2000, _is_self_play=0):
        self.mcts = MCTS_candq(policy_value_function, c_1, c_2, n_playout, _is_self_play=_is_self_play)
        self._is_self_play = _is_self_play

    # def set_player_ind(self, p):
    #     self.player = p

    def reset_sparse(self):
        self.mcts.update_with_move(-1)

    def get_action(self, board, temp=1e-3, return_prob=0):
        sensible_moves = board.availables
        # the pi vector returned by MCTS as in the alphaGo Zero paper
        move_probs = np.zeros(board.height * 32)
        
        if len(sensible_moves) > 0:
            acts, probs = self.mcts.get_move_probs(board, temp)
            move_probs[list(acts)] = probs
            
            # print(move_probs)
            
            if self._is_self_play:
                # add Dirichlet Noise for exploration (needed for self-play training)
                move = np.random.choice(
                    acts,
                    p=0.8*probs + 0.2*np.random.dirichlet(0.08*np.ones(len(probs)))
                    # p=probs
                )
                # update the root node and reuse the search tree
                self.mcts.update_with_move(move)
            else:
                # with the default temp=1e-3, it is almost equivalent
                # to choosing the move with the highest prob
                move = np.random.choice(acts, p=probs)
                # reset the root node
                self.mcts.update_with_move(-1)
#                location = board.move_to_location(move)
#                print("AI move: %d,%d\n" % (location[0], location[1]))

            move_policy = move // 32
            move_policy_value = move % 32
            
            move_policy_prob = np.sum(move_probs.reshape((board.height, 32)), axis = -1)
            move_policy_value_prob = move_probs[move_policy*32 : move_policy*32 + 32] / move_policy_prob[move_policy]

            if return_prob:
                return move_policy, move_policy_prob, move_policy_value, move_policy_value_prob
                # return move, move_probs
            else:
                return move_policy, move_policy_value, move_policy_prob, move_policy_value_prob
        else:
            print("WARNING: the board is full")


#     def get_action_eval(self, board, temp=1e-3, return_prob=0):
#         sensible_moves = board.availables
#         # the pi vector returned by MCTS as in the alphaGo Zero paper
#         move_probs = np.zeros(board.height * 32)
        
#         if len(sensible_moves) > 0:
#             acts, probs = self.mcts.get_move_probs(board, temp)
#             move_probs[list(acts)] = probs
            
#             if self._is_selfplay:
#                 # add Dirichlet Noise for exploration (needed for
#                 # self-play training)
#                 move = np.random.choice(
#                     acts,
#                     p=0.75*probs + 0.25*np.random.dirichlet(0.3*np.ones(len(probs)))
#                 )
#                 # update the root node and reuse the search tree
#                 self.mcts.update_with_move(move)
#             else:
#                 # with the default temp=1e-3, it is almost equivalent
#                 # to choosing the move with the highest prob
#                 move = np.random.choice(acts, p=probs)
#                 # reset the root node
#                 self.mcts.update_with_move(-1)
# #                location = board.move_to_location(move)
# #                print("AI move: %d,%d\n" % (location[0], location[1]))


#             move_policy = move // 32
#             move_policy_value = move % 32
            
#             move_policy_prob = np.sum(move_probs.reshape((board.height, 32)), axis = -1)
#             move_policy_value_prob = move_probs[move_policy*32 : move_policy*32 + 32] / move_policy_prob[move_policy]

#             if return_prob:
#                 return move_policy, move_policy_prob, move_policy_value, move_policy_value_prob
#                 # return move, move_probs
#             else:
#                 return move_policy, move_policy_value
#         else:
#             print("WARNING: the board is full")
    

    # def __str__(self):
    #     return "MCTS {}".format(self.player)
